"""
===============================
@Author     : Zuo WenTao
@Time       : 2024/10/30 10:58
@Description: 
@Change     : 
@File       : DownloadEmailAttachment.py
===============================
"""
import datetime
import os
import poplib
import re
from email import policy
from email.message import Message
from email.parser import Parser, BytesParser
from email.utils import parseaddr
from nntplib import decode_header

import fitz
from loguru import logger

from Api.BatchAttachmentDownloader.emailinfo import EmailInfo, DateJudge


class NotReceiverError(Exception): ...


class DownloadEmailAttachment:

    def __init__(self, email_server, email_address, email_password, save_dir):
        self.email_address = email_address
        self.save_dir = save_dir
        self.date_begin = '2024-10-30 00:00'
        self.time_zone = "+0800"
        self.seve_file_dict = {}
        self.error_file_dist = {}
        self.__receiver = Pop3Receiver(email_server, email_address, email_password)

    def decode_str(self, s):
        """解码邮件头部的字符串"""
        decoded, charset = decode_header(s)[0]
        if charset:
            decoded = decoded.decode(charset)
        return decoded

    def verify_pdf(self, file_path):
        try:
            Doc = fitz.open(file_path)
            page = Doc.page_count
            if not page:
                return False
            return True
        except:
            return False

    def __save_attachment(self, msg, subject, directory):

        for part in msg.iter_attachments():
            filename = part.get_filename()
            if filename in directory:
                continue
            if filename:
                filename = filename if isinstance(filename, str) else self.decode_str(filename)

                file_data = part.get_payload(decode=True)

                file_path = os.path.join(self.save_dir, filename)
                with open(file_path, 'wb') as f:
                    f.write(file_data)

                # if self.verify_pdf(file_path):
                self.seve_file_dict[subject] = file_path
                logger.success(f"附件已保存: {file_path}")
                # else:
                #     os.remove(file_path)
                #     logger.warning(f"附件已删除: {file_path}")
                #     self.error_file_dist[subject] = file_path

    def download_attachments(self,directory):
        if self.__receiver is None:
            raise NotReceiverError(f'{self.email_address} -- 链接邮箱失败，请检查邮箱地址或网络连接。')
        # mail_quantity, mail_total_size = self.__receiver.get_email_status()
        # print('邮件总数:', mail_quantity)
        # if mail_total_size > 0:
        #     logger.debug('邮件总大小:', EmailInfo.bytes_to_readable(mail_total_size), end='\n\n')

        mail_list = self.__receiver.get_mail_list()
        error_count = 0
        for mail_number in mail_list:
            try:
                content_byte = self.__receiver.get_mail_header_bytes(mail_number)
                mail_message = self.parse_mail_byte_content(content_byte)
                message_info = self.__get_email_info(mail_message)
            except Exception as e:
                print('邮件接收或解码失败，邮件编号：[%s]  错误信息：%s' % (mail_number, e))
                error_count += 1
                continue

            # 超出设定的最早时间则结束循环
            if DateJudge.is_earlier(message_info.date, self.date_begin + self.time_zone):
                break

            response, lines, octets = self.__receiver.retr(mail_number)
            msg_content = b'\r\n'.join(lines)

            msg = BytesParser(policy=policy.default).parsebytes(msg_content)
            subject = msg['subject'].replace("回复：","") if isinstance(msg['subject'], str) else self.decode_str(msg['subject']).replace("回复：","")

            self.__save_attachment(msg, subject, directory)

        return self.seve_file_dict, self.error_file_dist

    @staticmethod
    def parse_mail_byte_content(content_byte):
        try:
            mail_content = content_byte.decode()
        except UnicodeDecodeError as e:
            mail_content = content_byte.decode(encoding='GB18030', errors='replace')

        return Parser().parsestr(mail_content)

    def __get_email_info(self, message: Message):
        email_info = EmailInfo()

        try:
            email_info.subject = self.decode_mail_info_str(message.get('Subject'))
        except TypeError as e:
            email_info.subject = '无主题'

        name, address = parseaddr(message.get('From'))
        email_info.from_address = address
        email_info.from_name = self.decode_mail_info_str(name)
        email_info.to_names, email_info.to_addresses = self.__parse_mail_reciver_info(
            message.get_all("To")
        )

        date = message.get('Date')
        # 少数情况下无Data字段，尝试从其他字段中获取时间
        if date is None:
            if date is None:
                date = decode_time_from_received(message)
            if date is None:
                date = decode_time_from_x_qq_mid(message)
            # 如果还不能获取到Date则报错。极少数邮件信息头内没有时间信息，偶发于一些系统发送的邮件
            if date is None:
                raise ValueError('该邮件收件时间解析失败，邮件主题：【%s】' % email_info.subject)

        # Date格式统一，最终格式：'4 Jan 2020 11:59:25 +0800' (%d %b %Y %H:%M:%S %z)
        # 通常源数据格式为'Sat, 4 Jan 2020 11:59:25 +0800', 也有可能是'4 Jul 2019 21:37:08 +0800'
        # 开头星期去除，部分数据末尾有附加信息，因此以首个冒号后+12截取
        date_begin_index = 0
        for date_begin_index in range(len(date)):
            if '0' <= date[date_begin_index] <= '9':
                break
        date = date.replace('GMT', '+0000')  # 部分邮件用GMT表示
        date = date[date_begin_index:date.find(':') + 12]
        # 最后以时间戳保存
        email_info.date = datetime.datetime.strptime(date, '%d %b %Y %H:%M:%S %z').timestamp()

        return email_info

    @staticmethod
    # 将邮件中的bytes数据转为字符串
    def decode_mail_info_str(content):
        result_content = []
        decode_value = decode_header(content)
        if isinstance(decode_value, str):
            return decode_value
        for value, charset in decode_value:
            if type(value) != str:
                if charset is None:
                    value = value.decode(errors='replace')
                elif charset.lower() in ['gbk', 'gb2312', 'gb18030']:
                    value = value.decode(encoding='gb18030', errors='replace')
                else:
                    value = value.decode(charset, errors='replace')

            result_content.append(value)
        return ' '.join(result_content)

    # 解析收件人地址名称
    def __parse_mail_reciver_info(self, to_address_list):
        to_names = list()
        to_addresses = list()
        if to_address_list[0].split(","):
            for address in to_address_list[0].split(","):
                name, email = parseaddr(address)
                decoded_string = self.decode_mail_info_str(name)
                if decoded_string:
                    to_names.append(decoded_string)
                if email:
                    to_addresses.append(email)
        return to_names, to_addresses


def decode_time_from_x_qq_mid(message: Message):
    field_str = message.get('X-QQ-mid')
    if field_str is None:
        return None

    time_str = re.findall('t[0-9]{10}t', field_str)
    if len(time_str) == 0:
        return None
    time = datetime.datetime.fromtimestamp(float(time_str[0][1:-1]), datetime.timezone.utc)
    return time.strftime('%d %b %Y %H:%M:%S %z')


def decode_time_from_received(message: Message):
    field_str = message.get('Received')
    if field_str is None:
        return None
    return field_str[field_str.rfind(';') + 1:]


class Pop3Receiver:
    def __init__(self, host: str, email_address: str, email_password: str):
        try:
            self.__connection = poplib.POP3_SSL(host)
        except OSError as e:
            logger.debug('连接服务器失败，请检查服务器地址或网络连接。')
            self.close()
            return

        self.__connection.set_debuglevel(False)
        poplib._MAXLINE = 65356

        logger.debug(self.__connection.getwelcome().decode())

        self.__connection.user(email_address)
        try:
            self.__connection.pass_(email_password)
        except Exception as e:
            logger.debug(e.args)
            logger.debug('登陆失败，请检查用户名/密码。并确保您的邮箱已开启POP3服务。')
            self.close()
            return

    def get_mail_list(self, condition: dict = None):
        response, mail_list, octets = self.__connection.list()
        return [x.split()[0].decode() for x in reversed(mail_list)]

    def get_email_status(self):
        return self.__connection.stat()

    def get_mail_header_bytes(self, mail_number: str):
        response, content_byte, octets = self.__connection.top(mail_number, 40)
        try:
            mail_header_end = content_byte.index(b'')
        except ValueError as e:
            mail_header_end = len(content_byte)
        return self.__merge_bytes_list(content_byte[:mail_header_end])

    def get_full_mail_bytes(self, mail_number: str):
        response, content_byte, size = self.__connection.retr(mail_number)
        return self.__merge_bytes_list(content_byte), size

    @staticmethod
    def __merge_bytes_list(bytes_list):
        return b'\n'.join(bytes_list)

    def retr(self, mail_number: str):
        return self.__connection.retr(mail_number)

    def close(self):
        if self.__connection is not None:
            try:
                self.__connection.close()
            except OSError as e:
                logger.debug('断开时发生错误')
            self.__connection = None
